#!/usr/bin/python
# -*-coding:utf-8-*-
import os
from functools import wraps
from time import time
from copy import deepcopy
import _pickle as cPickle

import pandas as pd
import numpy as np

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

data_reader = DataReader()

def d_type_exam(d, default=1.0):
    if isinstance(d, (float, int, np.float32)):
        return d
    else:
        print('d type error:', type(d),'\nd:', d, '\nreturned:', default)
        return default

def d_map(d, base=5):
    # d = 0.002
    try:
        str_d = str(d)
        str_d = str_d.replace('0', '')
        str_d = str_d.replace('.', '')

        return int(str_d[0]) * base
    except:
        return int(1.0 * base)

def d_map_limit(d, base=5, lowlimit=1, uplimit=3):
    # d = 0.002
    try:
        str_d = str(d)
        str_d = str_d.replace('0', '')
        str_d = str_d.replace('.', '')

        multi_d = int(str_d[0])

        if multi_d > uplimit:
            multi_d = uplimit
        elif multi_d < lowlimit:
            multi_d = lowlimit
        else:
            multi_d = multi_d

        return  multi_d * base
    except:
        return int(1.0 * base)


def disperse_map(v, y, x=None):
    """
    :param v: input random seed
    :param y: y choose spread
    :param x: x is v spread
    :return: v associated y
    y = [111,222,333]
    x = [0,30,70,100]
    v = 8
    """

    if x is None:
        x=np.linspace(1,100,len(y)+1)

    for i in range(len(y)):
        if v>=x[i] and v<=x[i+1]:
            return y[i]

# args传入list of np.array,然后给np.array添加keys
# def keys_core_pd(keys, *args):
#     data = deepcopy(keys)
#     for num, feature in enumerate(args):
#         try:
#             data['factor_'+str(num+1)]=feature
#         except:
#             print('keys:',keys)
#             print('feature:',feature)
#             print('data:',data)
#     return data

def np_keys_to_pd(keys):
    keys = pd.DataFrame(keys,
                        columns=['date', 'stock_code']).set_index(['date', 'stock_code']).index

    return keys

def np_neu_keys_to_pd(neu_keys):
    neu_keys = pd.DataFrame(neu_keys,
                            columns=['date', 'stock_code', 'ci1_code', 'cap'])

    neu_keys = neu_keys.set_index(['date', 'stock_code'])

    return neu_keys

def np_filter_season_keys_to_pd(filter_season_keys):
    filter_season_keys = pd.DataFrame(filter_season_keys,
                                      columns=['date', 'stock_code', 'drop'])

    filter_season_keys = filter_season_keys.set_index(['date', 'stock_code'])

    return filter_season_keys

def keys_core_pd(keys, *args):
    data = pd.DataFrame(index=keys)
    for num, feature in enumerate(args):
        data['f%s' % (num+1)] = feature

    return data

def keys_core_pd_v2(keys, *args):
    # keys = pd.DataFrame(keys, columns=['date', 'stock_code']).set_index(['date', 'stock_code']).index
    keys = np_keys_to_pd(keys)
    data = pd.DataFrame(index=keys)

    for num, feature in enumerate(args):
        data['f%s' % (num+1)] = feature

    return data

# 获取最近的未来n个交易日
def get_next_trade_date(end_date, n=1):
    trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    trade_date_data = pd.DatetimeIndex(trade_date_data['trade_date'].unique())

    trade_date_data = trade_date_data[trade_date_data >= end_date]

    return trade_date_data[n-1].strftime('%Y-%m-%d')

def get_previous_trade_date(end_date, n=1):
    data_reader = DataReader()

    trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    trade_date_data = pd.DatetimeIndex(trade_date_data['trade_date'].unique())

    trade_date_data = trade_date_data[trade_date_data <= end_date]

    return trade_date_data[-n].strftime('%Y-%m-%d')

def get_trade_days(start_date, end_date):
    trade_date = data_reader.read_basic_data_table('processed_trade_date_data')
    trade_date = pd.DatetimeIndex(trade_date['trade_date'].unique())

    trade_date = trade_date[(trade_date >= start_date) & (trade_date <= end_date)]

    return trade_date.strftime('%Y-%m-%d')

def save_txt_file(path, data):
    with open(path, 'w') as f:
        f.write(data)

def delete_file_list(del_filepath_list):
    for file in del_filepath_list:
        try:
            os.remove(file)
            print(file, 'remove done!')
        except:
            print(file, 'remove exception!')


def writeToFile(file, store_path, mode='wb'):
    with open(store_path, mode) as f:
        cPickle.dump(file, f)

def loadFile(file_path, mode='rb'):
    with open(file_path, mode) as f:
        return cPickle.load(f)


